import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):  # 建议使用einsum实现
    def __init__(self, heads, d_model, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_k = d_model // heads
        self.h = heads

        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask):
        bs = q.size(0)

        # 线性变换并为多头注意力准备
        k = self.k_linear(k).view(bs, -1, self.h, self.d_k)  # -1是自动计算len
        q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
        v = self.v_linear(v).view(bs, -1, self.h, self.d_k)

        # 维度调整
        k = k.transpose(1, 2)  # (B,H,L,d_k)
        q = q.transpose(1, 2)
        v = v.transpose(1, 2)

        # 注意力计算
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            if q.shape[-2] != mask.shape[-2]:
                mask = F.pad(mask, (0, 0, 0, 1), "constant", 0)
            scores = scores.masked_fill(mask == 0, value=-1e9)  # 假设 mask 是一个布尔张量
        attention = F.softmax(scores, dim=-1)
        attention = self.dropout(attention)

        # 应用注意力
        '''
        attention 张量 [n,h,q,l]:
            n: batch size
            h: number of heads（注意力头数）
            q: query sequence length
            l: key/value sequence length
        
        values 张量 [n,l,h,d]:
            n: batch size
            l: key/value sequence length
            h: number of heads
            d: head dimension（每个头的维度大小）
        
        输出张量 [n,q,h,d]:
            n: batch size
            q: query sequence length
            h: number of heads
            d: head dimension
        '''
        x = torch.matmul(attention, v)
        x = x.transpose(1, 2).contiguous().view(bs, -1, self.d_model)  # 建议使用einsum实现
        x = self.out(x)
        return x


class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = F.relu(self.linear_1(x))
        x = self.dropout(x)
        x = self.linear_2(x)
        return x


class EncoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()
        self.norm_1 = nn.LayerNorm(d_model)
        self.norm_2 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.ff = FeedForward(d_model, dropout=dropout)
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x, mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn(x2, x2, x2, mask))  # 注意力机制与残差连接
        x2 = self.norm_2(x)  # layerNorm
        x = x + self.dropout_2(self.ff(x2))  # MLP 与残差连接
        return x


class DecoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()
        self.norm_1 = nn.LayerNorm(d_model)
        self.norm_2 = nn.LayerNorm(d_model)
        self.norm_3 = nn.LayerNorm(d_model)
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)
        self.dropout_3 = nn.Dropout(dropout)

        self.attn_1 = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.attn_2 = MultiHeadAttention(heads, d_model, dropout=dropout)
        self.ff = FeedForward(d_model, dropout=dropout)

    def forward(self, x, e_outputs, src_mask, trg_mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn_1(x2, x2, x2, trg_mask))  # 解码器自注意力
        x2 = self.norm_2(x)
        x = x + self.dropout_2(self.attn_2(x2, e_outputs, e_outputs, src_mask))  # 编码器-解码器注意力
        x2 = self.norm_3(x)
        x = x + self.dropout_3(self.ff(x2))  # 前馈网络
        return x


class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads):
        super().__init__()
        self.encoder = Encoder(vocab_size, d_model, N, heads)
        self.decoder = Decoder(vocab_size, d_model, N, heads)
        self.out = nn.Linear(d_model, vocab_size)
        self.softmax = nn.LogSoftmax(dim=-1)  # 最终的softmax操作

    def forward(self, src, trg, src_mask, trg_mask):
        e_outputs = self.encoder(src, src_mask)
        d_output = self.decoder(trg, e_outputs, src_mask, trg_mask)
        output = self.out(d_output)
        return self.softmax(output)


class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)  # nn.Embedding 是用来做词嵌入(word embedding)
        self.pe = PositionalEncoder(d_model)
        # DecoderLayer 是 Transformer 模型中的一个核心组件，它构成了解码器部分的基本单元。
        self.layers = nn.ModuleList([EncoderLayer(d_model, heads) for _ in range(N)])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, src, mask):
        x = self.embed(src)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, N, heads):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model)
        self.layers = nn.ModuleList([DecoderLayer(d_model, heads) for _ in range(N)])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, trg, e_outputs, src_mask, trg_mask):
        x = self.embed(trg)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, e_outputs, src_mask, trg_mask)
        return self.norm(x)


class PositionalEncoder(nn.Module):
    def __init__(self, d_model, max_seq_len=200):
        super().__init__()
        self.d_model = d_model
        pe = torch.zeros(max_seq_len, d_model)
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i) / d_model)))
                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1)) / d_model)))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x * math.sqrt(self.d_model)
        x = x + self.pe[:, :x.size(1)]
        return x


if __name__ == '__main__':
    # 示例使用
    vocab_size = 10000
    d_model = 512
    N = 2
    heads = 8
    src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
    trg = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])
    src_mask = None  # 编码器使用的，让编码器能够看到整个输入序列
    trg_mask = torch.zeros(11, 11)
    model = Transformer(vocab_size, d_model, N, heads)
    out = model(src, trg, src_mask, trg_mask)
    print(out)  # 输出结果
